import sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from scipy.linalg import solve_continuous_are
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QGroupBox, QLabel, QLineEdit, QPushButton, QTabWidget, QGridLayout,
                             QDoubleSpinBox, QSpinBox, QComboBox, QTextEdit, QSplitter)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
import time
import logging
from io import StringIO

# 设置日志
log_stream = StringIO()
logging.basicConfig(stream=log_stream, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger('DroneLQR')


class SimulationWorker(QThread):
    """仿真工作线程"""
    finished = pyqtSignal(dict)
    progress = pyqtSignal(int)
    log = pyqtSignal(str)

    def __init__(self, params):
        super().__init__()
        self.params = params
        self.is_running = True

    def run(self):
        """运行仿真"""
        try:
            # 解包参数
            m = self.params['mass']
            g = self.params['gravity']
            Ixx, Iyy, Izz = self.params['inertia']
            dt = self.params['dt']
            total_time = self.params['total_time']
            Q_weights = self.params['Q_weights']
            R_weights = self.params['R_weights']
            initial_state = self.params['initial_state']
            target_state = self.params['target_state']

            steps = int(total_time / dt)

            # 初始化状态
            state = np.array(initial_state, dtype=np.float64)
            target = np.array(target_state, dtype=np.float64)

            # 计算LQR增益
            A = np.zeros((12, 12), dtype=np.float64)
            A[0:3, 6:9] = np.eye(3)
            A[3:6, 9:12] = np.eye(3)
            A[6, 4] = g
            A[7, 3] = -g

            B = np.zeros((12, 4), dtype=np.float64)
            B[8, 0] = 1 / m
            B[9:12, 1:4] = np.array([[1 / Ixx, 0, 0],
                                     [0, 1 / Iyy, 0],
                                     [0, 0, 1 / Izz]], dtype=np.float64)

            Q = np.diag(Q_weights)
            R = np.diag(R_weights)

            P = solve_continuous_are(A, B, Q, R)
            K = np.linalg.inv(R) @ B.T @ P

            # 存储历史数据
            history = {
                'time': np.zeros(steps, dtype=np.float64),
                'position': np.zeros((steps, 3), dtype=np.float64),
                'attitude': np.zeros((steps, 3), dtype=np.float64),
                'velocity': np.zeros((steps, 3), dtype=np.float64),
                'angular_velocity': np.zeros((steps, 3), dtype=np.float64),
                'control': np.zeros((steps, 4), dtype=np.float64),
                'error': np.zeros((steps, 3), dtype=np.float64)
            }

            # 运行仿真
            for i in range(steps):
                if not self.is_running:
                    break

                # 计算控制输入
                error = state - target
                u = -K @ error
                u[0] = max(u[0], 0)  # 确保推力非负

                # 无人机动力学
                x, y, z, phi, theta, psi, vx, vy, vz, p, q, r = state
                F, tau_x, tau_y, tau_z = u

                # 位置导数
                dx = vx
                dy = vy
                dz = vz

                # 欧拉角导数
                if abs(np.cos(theta)) < 1e-5:
                    theta = np.sign(theta) * (np.pi / 2 - 1e-5)

                dphi = p + (q * np.sin(phi) + r * np.cos(phi)) * np.tan(theta)
                dtheta = q * np.cos(phi) - r * np.sin(phi)
                dpsi = (q * np.sin(phi) + r * np.cos(phi)) / np.cos(theta)

                # 速度导数
                dvx = -F / m * (np.sin(phi) * np.sin(psi) + np.cos(phi) * np.cos(psi) * np.sin(theta))
                dvy = F / m * (np.cos(psi) * np.sin(phi) - np.cos(phi) * np.sin(psi) * np.sin(theta))
                dvz = F / m * (np.cos(phi) * np.cos(theta)) - g

                # 角速度导数
                dp = (tau_x + (Iyy - Izz) * q * r) / Ixx
                dq = (tau_y + (Izz - Ixx) * p * r) / Iyy
                dr = (tau_z + (Ixx - Iyy) * p * q) / Izz

                derivatives = np.array([dx, dy, dz, dphi, dtheta, dpsi, dvx, dvy, dvz, dp, dq, dr])

                # RK4积分
                k1 = derivatives
                k2 = self.dynamics(state + 0.5 * dt * k1, u, m, g, Ixx, Iyy, Izz)
                k3 = self.dynamics(state + 0.5 * dt * k2, u, m, g, Ixx, Iyy, Izz)
                k4 = self.dynamics(state + dt * k3, u, m, g, Ixx, Iyy, Izz)

                state += dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6

                # 存储数据
                history['time'][i] = i * dt
                history['position'][i] = state[0:3]
                history['attitude'][i] = state[3:6]
                history['velocity'][i] = state[6:9]
                history['angular_velocity'][i] = state[9:12]
                history['control'][i] = u
                history['error'][i] = state[0:3] - target[0:3]

                # 更新进度
                if i % 100 == 0:
                    self.progress.emit(int(i / steps * 100))
                    pos_error = np.linalg.norm(state[0:3] - target[0:3])
                    self.log.emit(f"Time: {i * dt:.2f}s, Position Error: {pos_error:.4f}m")

            self.finished.emit(history)

        except Exception as e:
            self.log.emit(f"Simulation error: {str(e)}")

    def dynamics(self, state, u, m, g, Ixx, Iyy, Izz):
        """计算状态导数"""
        x, y, z, phi, theta, psi, vx, vy, vz, p, q, r = state
        F, tau_x, tau_y, tau_z = u

        # 位置导数
        dx = vx
        dy = vy
        dz = vz

        # 避免角度奇异点
        if abs(np.cos(theta)) < 1e-5:
            theta = np.sign(theta) * (np.pi / 2 - 1e-5)

        # 欧拉角导数
        dphi = p + (q * np.sin(phi) + r * np.cos(phi)) * np.tan(theta)
        dtheta = q * np.cos(phi) - r * np.sin(phi)
        dpsi = (q * np.sin(phi) + r * np.cos(phi)) / np.cos(theta)

        # 速度导数
        dvx = -F / m * (np.sin(phi) * np.sin(psi) + np.cos(phi) * np.cos(psi) * np.sin(theta))
        dvy = F / m * (np.cos(psi) * np.sin(phi) - np.cos(phi) * np.sin(psi) * np.sin(theta))
        dvz = F / m * (np.cos(phi) * np.cos(theta)) - g

        # 角速度导数
        dp = (tau_x + (Iyy - Izz) * q * r) / Ixx
        dq = (tau_y + (Izz - Ixx) * p * r) / Iyy
        dr = (tau_z + (Ixx - Iyy) * p * q) / Izz

        return np.array([dx, dy, dz, dphi, dtheta, dpsi, dvx, dvy, dvz, dp, dq, dr])

    def stop(self):
        """停止仿真"""
        self.is_running = False


class MplCanvas(FigureCanvas):
    """Matplotlib画布"""

    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)

        self.ax = self.fig.add_subplot(111)

    def plot(self, x, y, labels, title, xlabel, ylabel):
        """绘制数据"""
        self.ax.clear()
        for i in range(len(y)):
            self.ax.plot(x, y[i], label=labels[i])
        self.ax.set_title(title)
        self.ax.set_xlabel(xlabel)
        self.ax.set_ylabel(ylabel)
        self.ax.legend()
        self.ax.grid(True)
        self.draw()


class ThreeDCanvas(FigureCanvas):
    """3D Matplotlib画布"""

    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)

        self.ax = self.fig.add_subplot(111, projection='3d')

    def plot_trajectory(self, x, y, z, target):
        """绘制3D轨迹"""
        self.ax.clear()
        self.ax.plot(x, y, z, 'b-', label='Trajectory')
        self.ax.scatter(target[0], target[1], target[2], c='red', marker='o', s=100, label='Target')
        self.ax.scatter(x[0], y[0], z[0], c='green', marker='x', s=100, label='Start')
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')
        self.ax.set_zlabel('Z (m)')
        self.ax.set_title('3D Trajectory')
        self.ax.legend()
        self.ax.grid(True)
        self.draw()


class MainWindow(QMainWindow):
    """主窗口"""

    def __init__(self):
        super().__init__()
        self.setWindowTitle("LQR无人机控制仿真")
        self.setGeometry(100, 100, 1200, 800)

        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QHBoxLayout(central_widget)

        # 左侧参数面板
        left_panel = QWidget()
        left_panel.setMaximumWidth(300)
        left_layout = QVBoxLayout(left_panel)

        # 参数组
        param_group = QGroupBox("仿真参数")
        param_layout = QGridLayout(param_group)

        # 物理参数
        param_layout.addWidget(QLabel("质量 (kg):"), 0, 0)
        self.mass_edit = QDoubleSpinBox()
        self.mass_edit.setRange(0.1, 10.0)
        self.mass_edit.setValue(1.0)
        self.mass_edit.setSingleStep(0.1)
        param_layout.addWidget(self.mass_edit, 0, 1)

        param_layout.addWidget(QLabel("重力 (m/s²):"), 1, 0)
        self.gravity_edit = QDoubleSpinBox()
        self.gravity_edit.setRange(0.1, 20.0)
        self.gravity_edit.setValue(9.81)
        self.gravity_edit.setSingleStep(0.1)
        param_layout.addWidget(self.gravity_edit, 1, 1)

        param_layout.addWidget(QLabel("Ixx (kg·m²):"), 2, 0)
        self.Ixx_edit = QDoubleSpinBox()
        self.Ixx_edit.setRange(0.01, 1.0)
        self.Ixx_edit.setValue(0.1)
        self.Ixx_edit.setSingleStep(0.01)
        param_layout.addWidget(self.Ixx_edit, 2, 1)

        param_layout.addWidget(QLabel("Iyy (kg·m²):"), 3, 0)
        self.Iyy_edit = QDoubleSpinBox()
        self.Iyy_edit.setRange(0.01, 1.0)
        self.Iyy_edit.setValue(0.1)
        self.Iyy_edit.setSingleStep(0.01)
        param_layout.addWidget(self.Iyy_edit, 3, 1)

        param_layout.addWidget(QLabel("Izz (kg·m²):"), 4, 0)
        self.Izz_edit = QDoubleSpinBox()
        self.Izz_edit.setRange(0.01, 1.0)
        self.Izz_edit.setValue(0.1)
        self.Izz_edit.setSingleStep(0.01)
        param_layout.addWidget(self.Izz_edit, 4, 1)

        # 仿真参数
        param_layout.addWidget(QLabel("时间步长 (s):"), 5, 0)
        self.dt_edit = QDoubleSpinBox()
        self.dt_edit.setRange(0.001, 0.1)
        self.dt_edit.setValue(0.01)
        self.dt_edit.setSingleStep(0.001)
        param_layout.addWidget(self.dt_edit, 5, 1)

        param_layout.addWidget(QLabel("总时间 (s):"), 6, 0)
        self.total_time_edit = QDoubleSpinBox()
        self.total_time_edit.setRange(1.0, 60.0)
        self.total_time_edit.setValue(10.0)
        self.total_time_edit.setSingleStep(1.0)
        param_layout.addWidget(self.total_time_edit, 6, 1)

        # 初始状态
        param_layout.addWidget(QLabel("初始高度 (m):"), 7, 0)
        self.init_z_edit = QDoubleSpinBox()
        self.init_z_edit.setRange(0.0, 10.0)
        self.init_z_edit.setValue(0.0)
        self.init_z_edit.setSingleStep(0.1)
        param_layout.addWidget(self.init_z_edit, 7, 1)

        param_layout.addWidget(QLabel("目标高度 (m):"), 8, 0)
        self.target_z_edit = QDoubleSpinBox()
        self.target_z_edit.setRange(0.0, 10.0)
        self.target_z_edit.setValue(5.0)
        self.target_z_edit.setSingleStep(0.1)
        param_layout.addWidget(self.target_z_edit, 8, 1)

        left_layout.addWidget(param_group)

        # LQR参数组
        lqr_group = QGroupBox("LQR参数")
        lqr_layout = QGridLayout(lqr_group)

        # Q矩阵权重
        q_labels = ["Q_x", "Q_y", "Q_z", "Q_φ", "Q_θ", "Q_ψ", "Q_vx", "Q_vy", "Q_vz", "Q_p", "Q_q", "Q_r"]
        self.q_edits = []
        default_q = [10, 10, 100, 10, 10, 10, 1, 1, 10, 1, 1, 1]

        for i, label in enumerate(q_labels):
            lqr_layout.addWidget(QLabel(label + ":"), i, 0)
            edit = QDoubleSpinBox()
            edit.setRange(0.1, 1000.0)
            edit.setValue(default_q[i])
            edit.setSingleStep(1.0)
            lqr_layout.addWidget(edit, i, 1)
            self.q_edits.append(edit)

        # R矩阵权重
        r_labels = ["R_F", "R_τx", "R_τy", "R_τz"]
        self.r_edits = []
        default_r = [0.1, 1, 1, 1]

        for i, label in enumerate(r_labels):
            lqr_layout.addWidget(QLabel(label + ":"), i, 2)
            edit = QDoubleSpinBox()
            edit.setRange(0.01, 10.0)
            edit.setValue(default_r[i])
            edit.setSingleStep(0.1)
            lqr_layout.addWidget(edit, i, 3)
            self.r_edits.append(edit)

        left_layout.addWidget(lqr_group)

        # 控制按钮
        self.run_button = QPushButton("开始仿真")
        self.run_button.clicked.connect(self.start_simulation)
        left_layout.addWidget(self.run_button)

        self.stop_button = QPushButton("停止仿真")
        self.stop_button.clicked.connect(self.stop_simulation)
        self.stop_button.setEnabled(False)
        left_layout.addWidget(self.stop_button)

        # 日志显示
        log_group = QGroupBox("日志")
        log_layout = QVBoxLayout(log_group)
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        log_layout.addWidget(self.log_text)
        left_layout.addWidget(log_group)

        # 右侧图形面板
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)

        # 标签页
        self.tabs = QTabWidget()

        # 位置标签页
        pos_tab = QWidget()
        pos_layout = QVBoxLayout(pos_tab)
        self.pos_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        pos_layout.addWidget(self.pos_canvas)
        pos_layout.addWidget(NavigationToolbar(self.pos_canvas, self))
        self.tabs.addTab(pos_tab, "位置")

        # 姿态标签页
        att_tab = QWidget()
        att_layout = QVBoxLayout(att_tab)
        self.att_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        att_layout.addWidget(self.att_canvas)
        att_layout.addWidget(NavigationToolbar(self.att_canvas, self))
        self.tabs.addTab(att_tab, "姿态")

        # 控制标签页
        ctrl_tab = QWidget()
        ctrl_layout = QVBoxLayout(ctrl_tab)
        self.ctrl_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        ctrl_layout.addWidget(self.ctrl_canvas)
        ctrl_layout.addWidget(NavigationToolbar(self.ctrl_canvas, self))
        self.tabs.addTab(ctrl_tab, "控制")

        # 3D轨迹标签页
        traj_tab = QWidget()
        traj_layout = QVBoxLayout(traj_tab)
        self.traj_canvas = ThreeDCanvas(self, width=5, height=4, dpi=100)
        traj_layout.addWidget(self.traj_canvas)
        traj_layout.addWidget(NavigationToolbar(self.traj_canvas, self))
        self.tabs.addTab(traj_tab, "3D轨迹")

        right_layout.addWidget(self.tabs)

        # 添加到主布局
        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([300, 900])
        main_layout.addWidget(splitter)

        # 仿真线程
        self.sim_thread = None

        # 初始日志
        self.log_text.append("欢迎使用LQR无人机控制仿真系统")
        self.log_text.append("请设置参数并点击'开始仿真'")

    def get_params(self):
        """获取参数"""
        params = {
            'mass': self.mass_edit.value(),
            'gravity': self.gravity_edit.value(),
            'inertia': [self.Ixx_edit.value(), self.Iyy_edit.value(), self.Izz_edit.value()],
            'dt': self.dt_edit.value(),
            'total_time': self.total_time_edit.value(),
            'Q_weights': [edit.value() for edit in self.q_edits],
            'R_weights': [edit.value() for edit in self.r_edits],
            'initial_state': [0.0, 0.0, self.init_z_edit.value()] + [0.0] * 9,
            'target_state': [0.0, 0.0, self.target_z_edit.value()] + [0.0] * 9
        }
        return params

    def start_simulation(self):
        """开始仿真"""
        self.log_text.append("开始仿真...")
        self.run_button.setEnabled(False)
        self.stop_button.setEnabled(True)

        params = self.get_params()
        self.sim_thread = SimulationWorker(params)
        self.sim_thread.finished.connect(self.simulation_finished)
        self.sim_thread.progress.connect(self.update_progress)
        self.sim_thread.log.connect(self.update_log)
        self.sim_thread.start()

    def stop_simulation(self):
        """停止仿真"""
        if self.sim_thread:
            self.sim_thread.stop()
            self.sim_thread.wait()
            self.log_text.append("仿真已停止")
            self.run_button.setEnabled(True)
            self.stop_button.setEnabled(False)

    def simulation_finished(self, history):
        """仿真完成"""
        self.log_text.append("仿真完成")
        self.run_button.setEnabled(True)
        self.stop_button.setEnabled(False)

        # 更新图表
        self.update_plots(history)

    def update_progress(self, value):
        """更新进度"""
        self.log_text.append(f"仿真进度: {value}%")

    def update_log(self, message):
        """更新日志"""
        self.log_text.append(message)

    def update_plots(self, history):
        """更新所有图表"""
        # 位置图表
        time = history['time']
        pos_x, pos_y, pos_z = history['position'][:, 0], history['position'][:, 1], history['position'][:, 2]
        self.pos_canvas.plot(time, [pos_x, pos_y, pos_z],
                             ['X', 'Y', 'Z'],
                             'Position vs Time', 'Time (s)', 'Position (m)')

        # 姿态图表
        att_phi, att_theta, att_psi = np.degrees(history['attitude'][:, 0]), np.degrees(
            history['attitude'][:, 1]), np.degrees(history['attitude'][:, 2])
        self.att_canvas.plot(time, [att_phi, att_theta, att_psi],
                             ['φ', 'θ', 'ψ'],
                             'Attitude vs Time', 'Time (s)', 'Angle (degrees)')

        # 控制图表
        thrust = history['control'][:, 0]
        torque_x, torque_y, torque_z = history['control'][:, 1], history['control'][:, 2], history['control'][:, 3]
        self.ctrl_canvas.plot(time, [thrust, torque_x, torque_y, torque_z],
                              ['Thrust', 'Torque X', 'Torque Y', 'Torque Z'],
                              'Control Inputs vs Time', 'Time (s)', 'Input Value')

        # 3D轨迹
        self.traj_canvas.plot_trajectory(pos_x, pos_y, pos_z, [0, 0, self.target_z_edit.value()])


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())